"""
Adapter to connect the barotropic model with the variational data assimilation framework
"""
import numpy as np
from Tools.barotropic_model import BarotropicModel
from Tools.AssimilationSystem import ObservationOperator

class BarotropicModelAdapter:
    """Adapter to make BarotropicModel compatible with Var4D requirements"""
    def __init__(self, model):
        self.model = model
        self.M = model.nx * model.ny  # State dimension
        
    def linear(self, x, t, dt):
        """Tangent linear model approximation"""
        # For nonlinear models, use finite difference approximation
        eps = 1e-8
        M = self.M
        TL = np.zeros((M, M))
        
        # Create perturbed states
        x = x.reshape(self.model.shape)
        for i in range(M):
            dx = np.zeros(M)
            dx[i] = eps
            perturbed = x.flatten() + dx
            TL[:,i] = (self.model.step(perturbed, dt) - self.model.step(x.flatten(), dt))/eps
            
        return TL

class BarotropicObservationOperator(ObservationOperator):
    """
    Specialized observation operator for barotropic model
    Can implement specific observation types like wind observations
    """
    def __init__(self, model, state_size, obs_size, obs_type='vorticity', obs_indices=None, seed=None):
        """
        Initialize barotropic observation operator
        
        Args:
            model: BarotropicModel instance
            state_size: Size of state vector
            obs_size: Number of observations
            obs_type: Type of observations ('vorticity', 'wind', 'streamfunction')
            obs_indices: Specific indices to observe (if None, randomly selected)
        """
        super().__init__(state_size, obs_size, obs_indices, seed=seed)
        self.model = model
        self.obs_type = obs_type
        
    def h(self, x):
        """Apply observation operator to state vector"""
        if self.obs_type == 'vorticity':
            # Direct observation of vorticity at selected points
            return super().h(x)
        
        elif self.obs_type == 'wind':
            # Wind observations require computing streamfunction and derivatives
            vorticity_2d = self.model._reshape_state(x)
            psi = self.model._compute_streamfunction(vorticity_2d)
            u, v = self.model._compute_velocity(psi)
            
            # Flatten and select points for u and v
            uv = np.concatenate([u.flatten(), v.flatten()])
            wind_obs = np.zeros(self.obs_size)
            
            # Half the observations are u, half are v
            half_obs = self.obs_size // 2
            for i in range(half_obs):
                wind_obs[i] = uv[self.obs_indices[i]]
                wind_obs[i + half_obs] = uv[self.model.state_size + self.obs_indices[i]]
                
            return wind_obs
            
        elif self.obs_type == 'streamfunction':
            # Streamfunction observations
            vorticity_2d = self.model._reshape_state(x)
            psi = self.model._compute_streamfunction(vorticity_2d)
            psi_flat = psi.flatten()
            
            return psi_flat[self.obs_indices]
            
        else:
            raise ValueError(f"Unknown observation type: {self.obs_type}")
            
    def h_linearized(self, x_ref):
        """
        Linearized observation operator
        For complex observation types, use finite differences
        """
        if self.obs_type == 'vorticity':
            return super().h_linearized(x_ref)
        else:
            # Use finite differences for more complex observation operators
            eps = 1e-6
            H = np.zeros((self.obs_size, self.state_size))
            
            for i in range(self.state_size):
                dx = np.zeros(self.state_size)
                dx[i] = eps
                
                y_plus = self.h(x_ref + dx)
                y_minus = self.h(x_ref)
                
                H[:, i] = (y_plus - y_minus) / eps
                
            return H
